from enum import Enum, StrEnum
from typing import Any, Literal, Union

from pydantic import BaseModel

from core.prompt.entities.advanced_prompt_entities import MemoryConfig
from core.tools.entities.tool_entities import ToolSelector
from core.workflow.nodes.base.entities import BaseNodeData


class AgentNodeData(BaseNodeData):
    agent_strategy_provider_name: str  # redundancy
    agent_strategy_name: str
    agent_strategy_label: str  # redundancy
    memory: MemoryConfig | None = None

    class AgentInput(BaseModel):
        value: Union[list[str], list[ToolSelector], Any]
        type: Literal["mixed", "variable", "constant"]

    agent_parameters: dict[str, AgentInput]


class ParamsAutoGenerated(Enum):
    CLOSE = 0
    OPEN = 1


class AgentOldVersionModelFeatures(StrEnum):
    """
    Enum class for old SDK version llm feature.
    """

    TOOL_CALL = "tool-call"
    MULTI_TOOL_CALL = "multi-tool-call"
    AGENT_THOUGHT = "agent-thought"
    VISION = "vision"
    STREAM_TOOL_CALL = "stream-tool-call"
    DOCUMENT = "document"
    VIDEO = "video"
    AUDIO = "audio"
